from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from contextlib import asynccontextmanager
from loguru import logger

from web.services.kafka_service import KafkaProducerService, KafkaConsumerService
from web.utils.logger_utils import setup_logger


# 定义 PDFRequest 模型（注意：callback_url 缺失，需根据业务补充）
class PDFRequest(BaseModel):
    doc_id: str
    callback_url: str  # 假设这是你实际需要的字段


# 定义 lifespan 管理应用生命周期
@asynccontextmanager
async def lifespan(this_app: FastAPI):
    # 初始化 Kafka 服务
    kafka_producer = KafkaProducerService()
    kafka_consumer = KafkaConsumerService()

    # 启动 Kafka 服务
    await kafka_producer.start()
    await kafka_consumer.start()
    logger.info("PDF Service started.")

    # 将服务挂载到 app.state 供路由使用
    this_app.state.kafka_producer = kafka_producer
    this_app.state.kafka_consumer = kafka_consumer

    yield  # 应用运行阶段

    # 关闭 Kafka 服务
    await kafka_producer.stop()
    await kafka_consumer.stop()
    logger.info("PDF Service stopped.")


# 创建 FastAPI 应用并绑定 lifespan
app = FastAPI(lifespan=lifespan)
setup_logger()


@app.post("/process-pdf")
async def process_pdf(request: PDFRequest):
    """Endpoint to submit PDF processing requests"""
    try:
        # 从 app.state 中获取 Kafka Producer
        producer = app.state.kafka_producer
        await producer.send_task(request.doc_id, request.callback_url)
        logger.info(f"Submitted PDF processing task for doc_id: {request.doc_id}")
        return {
            "status": "accepted",
            "doc_id": request.doc_id,
            "message": "PDF processing started. You will be notified via callback when complete."
        }
    except Exception as e:
        logger.error(f"Failed to submit task: {str(e)}")
        raise HTTPException(status_code=500, detail="Failed to submit processing task")

if __name__ == '__main__':
    print("dev command: uvicorn web.main:app --reload --host 0.0.0.0 --port 8000")
    print("pro command: gunicorn -w 4 -k uvicorn.workers.UvicornWorker -b 0.0.0.0:8000 web.main:app --preload")